import numpy as np
import matplotlib.pyplot as plt
from scipy.special import j0, j1, jn_zeros
from numpy.fft import fft, ifft, fftfreq

# ====================== 公共参数设置 ======================
a = 0.15       # 半径 (m)
peak = -0.7e6  # 应力峰值 (Pa)
T_load = 3 # Ft的周期 (s)
rho = 1500     # 密度 (kg/m3)
E = 100e6      # 弹性模量 (Pa)
nu = 0.35      # 泊松比

la = (nu * E)/((1 + nu) * (1 - 2 * nu))  # 拉梅常数
mu = E/(2 * (1 + nu))                    # 剪切模量
cp = np.sqrt((la + 2 * mu)/rho)          # 纵波波速
cs = np.sqrt(mu/rho)                     # 横波波速
omega0 = np.pi/T_load                    # 基频

# ====================== 时间参数 ======================
T_total = 61*T_load      # 总时长
St = 10*T_load           # 起始时间
Nt = 2048                # 采样点数
dt = T_total/Nt          # 时间步长
t = np.linspace(0, T_total, Nt)
mask = (t >= St) & (t <= St + T_load)
Ft = np.where(mask, peak*np.sin(np.pi*(t - St)/T_load), 0.0)

# ====================== 傅里叶变换 ======================
Ft_omega = fft(Ft)*dt
omega = 2*np.pi*fftfreq(Nt, dt)

# =================================================================
#                          方法1: 汉克尔变换
# =================================================================
def method1():
    # 波数参数
    k_max = 1000.0
    Nk = 6000
    k = np.linspace(1e-18, k_max, Nk)
    r_points = np.array([0, 0.2, 0.3, 0.45, 0.6, 0.9, 1.2, 1.5, 1.8])

    # 汉克尔变换核心计算
    def compute_Q(k):
        Q = np.zeros_like(k)
        mask = k != 0
        Q[mask] = (a * 1 / k[mask]) * j1(a * k[mask])
        Q[~mask] = 1 * a**2 / 2
        return Q

    Fr_k = compute_Q(k)
    sigma = np.outer(Fr_k, Ft_omega)

    def compute_B_D(k, omega):
        B = np.zeros((len(k), len(omega)), dtype=np.complex128)
        D = np.zeros_like(B, dtype=np.complex128)
        for i, omg in enumerate(omega):
            k_sq = k**2
            la_2mu = la + 2*mu
            omega_sq = omg**2
            if abs(omg) < 1e-30:
                B[:, i] = -1/(2*k*(la + mu))
                D[:, i] = la_2mu/(2*k*mu*(la + mu))
                continue
                
            sqrt_term1 = np.sqrt((k_sq*la_2mu - omega_sq*rho)/la_2mu + 0j)
            sqrt_term2 = np.sqrt((k_sq*mu - omega_sq*rho)/mu + 0j)
            denominator = 4*k_sq*sqrt_term1*sqrt_term2*mu**2 - 4*(k_sq*mu - 0.5*omega_sq*rho)**2
            B[:, i] = (k*(2*sqrt_term1*mu*sqrt_term2) - 2*k_sq*mu + omega_sq*rho)/denominator
            D[:, i] = (sqrt_term1*omega_sq*rho)/denominator
        return B, D

    B, D = compute_B_D(k, omega)
    w_k_omega = D * sigma

    # 逆变换
    J0_kr = j0(r_points[:, np.newaxis] * k)
    integrand = J0_kr[:, :, np.newaxis] * w_k_omega[np.newaxis, :, :] * k[np.newaxis, :, np.newaxis]
    w_r_omega = np.trapz(integrand, x=k, axis=1)
    w_r_t = np.real(ifft(w_r_omega/dt, axis=1))
    return w_r_t

# =================================================================
#                      方法2: 傅里叶-贝塞尔级数
# =================================================================
def method2():
    R = 25.0       # 有限域半径
    Nk = 6000   # 波数采样数
    alpha_n = jn_zeros(0, Nk)
    k = alpha_n/R
    J1_alpha = j1(alpha_n)
    r_points = np.array([0, 0.2, 0.3, 0.45, 0.6, 0.9, 1.2, 1.5, 1.8])

    # 傅里叶-贝塞尔系数
    Fr_k = np.zeros_like(k)
    for m in range(len(k)):
        alpha_m = alpha_n[m]
        scaled_alpha = alpha_m*a/R
        Fr_k[m] = 2*a*1*j1(scaled_alpha)/(alpha_m*R*J1_alpha[m]**2)

    # 核心计算
    def compute_B_D(k, omega):
        B = np.zeros((len(k), len(omega)), dtype=np.complex128)
        D = np.zeros_like(B, dtype=np.complex128)
        for i, omg in enumerate(omega):
            k_sq = k**2
            la_2mu = la + 2*mu
            omega_sq = omg**2
            if abs(omg) < 1e-30:
                B[:, i] = -1/(2*k*(la + mu))
                D[:, i] = la_2mu/(2*k*mu*(la + mu))
                continue
                
            sqrt_term1 = np.sqrt((k_sq*la_2mu - omega_sq*rho)/la_2mu + 0j)
            sqrt_term2 = np.sqrt((k_sq*mu - omega_sq*rho)/mu + 0j)
            denominator = 4*k_sq*sqrt_term1*sqrt_term2*mu**2 - 4*(k_sq*mu - 0.5*omega_sq*rho)**2
            B[:, i] = (k*(2*sqrt_term1*mu*sqrt_term2) - 2*k_sq*mu + omega_sq*rho)/denominator
            D[:, i] = (sqrt_term1*omega_sq*rho)/denominator
        return B, D

    B, D = compute_B_D(k, omega)
    sigma = np.outer(Fr_k, Ft_omega)
    w_k_omega = D * sigma

    # 逆变换
    J0_kr = j0(r_points[:, np.newaxis] * k)
    integrand = J0_kr[:, :, np.newaxis] * w_k_omega[np.newaxis, :, :]
    w_r_omega = np.sum(integrand, axis=1)
    w_r_t = np.real(ifft(w_r_omega/dt, axis=1))
    return w_r_t

# =================================================================
#                          结果可视化
# =================================================================
# 计算两种方法的结果
w1 = method1()
w2 = method2()
r_points = np.array([0, 0.2, 0.3, 0.45, 0.6, 0.9, 1.2, 1.5, 1.8])

# 计算所有数据的极值
all_values = np.concatenate([w1, w2])
abs_max = np.max(np.abs(all_values)) * 1.3
y_min = -abs_max
y_max = abs_max * 0.3

# 创建画布（增加高度到8英寸）
fig = plt.figure(figsize=(13, 6), dpi=300) 
colors = plt.cm.jet(np.linspace(0, 1, len(r_points)))

# ========================= 方法1 =========================
ax1 = plt.subplot2grid((5, 2), (0, 0), rowspan=3)
ax1_env = plt.subplot2grid((5, 2), (3, 0), rowspan=2)

# 绘图设置
for i, r_val in enumerate(r_points):
    ax1.plot(t, w1[i], color=colors[i], label=f'r={r_val:.2f}m', lw=1)
ax1.set_title(f'Hankel Transform ($T_d$={T_load}s)', pad=12)
ax1_env.plot(r_points, np.min(w1, axis=1), 'o-', color='navy', markersize=4)
for i, r_val in enumerate(r_points):
    w_val = np.min(w1[i])
    ax1_env.text(r_val -0.03, w_val - 0.0003, f'{w_val*1000:.2f}mm', fontsize=10)

# ========================= 方法2 =========================
ax2 = plt.subplot2grid((5, 2), (0, 1), rowspan=3)
ax2_env = plt.subplot2grid((5, 2), (3, 1), rowspan=2)

# 绘图设置
for i, r_val in enumerate(r_points):
    ax2.plot(t, w2[i], color=colors[i], label=f'r={r_val:.2f}m', lw=1)
ax2.set_title(f'Fourier-Bessel Series ($T_d$={T_load}s)', pad=12)
ax2_env.plot(r_points, np.min(w2, axis=1), 's-', color='maroon', markersize=4)
for i, r_val in enumerate(r_points):
    w_val = np.min(w2[i])
    ax2_env.text(r_val -0.03, w_val - 0.0003, f'{w_val*1000:.2f}mm', fontsize=10)

# ======================= 统一设置 ========================
for ax in [ax1, ax2]:
    ax.set_xlim(9.5*T_load, 12.5*T_load)
    ax.set_ylim(y_min, y_max)
    ax.grid(True, alpha=0.3)
    ax.set_ylabel('Displacement (m)', labelpad=-65)
    ax.set_xlabel('Time (s)', labelpad=-30)

for ax in [ax1_env, ax2_env]:
    ax.set_ylim(y_min, y_max)
    ax.grid(True, alpha=0.3)
    ax.set_ylabel('Max Deflection (m)', labelpad=-60)
    ax.set_xlabel('Radial Position (m)', labelpad=-30)

# ======================= 图例设置 ========================
handles, labels = ax1.get_legend_handles_labels()
fig.legend(handles, labels,
          loc='lower center',
          ncol=5,  # 减少列数
          borderaxespad=0.5,
          columnspacing=1.5,
          bbox_to_anchor=(0.5, 0.03),  # 下移图例
          frameon=False)

# ======================= 布局调整 ========================
plt.subplots_adjust(
    hspace=0.4,   # 增加行间距
    wspace=0.2,
    bottom=0.18,  # 增加底部空间
    top=0.92
)
plt.savefig(f'halfspace_displacement_response(T={T_load}s).png', bbox_inches='tight')
plt.show()